package com.it.smn.divide.loadbalance;

import com.it.smn.divide.datasource.AbstractDataSourceWrapper;
import com.it.smn.divide.datasource.ReadableDataSource;
import org.springframework.util.CollectionUtils;

import javax.sql.DataSource;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * @author sunmannian
 * @ClassName com.it.smn.divide.loadbalance.RoundRobinLoadBalance
 * @Date 2020/8/27 18:36
 * @Version V1.0
 * Description  <权重轮询负载均衡>
 * Copyright (C) 北京中电兴发科技有限公司  since 2020/8/27
 */
public class RoundRobinLoadBalance implements LoadBalanceStrategy {

    private WeightDataSource best;

    private List<WeightDataSource> weightDataSources;

    @Override
    public DataSource select(List<DataSource> dataSources) {
        return doSelect(dataSources);
    }

    private DataSource doSelect(List<DataSource> dataSources) {
        int size = dataSources.size();
        if (ONLY_ONE == size) {
            return dataSources.get(0);
        }
        if (CollectionUtils.isEmpty(weightDataSources)) {
            weightDataSources = convert(dataSources);
        }
        best = smooth(weightDataSources, best);
        return best.getDataSource();
    }

    /**
     * <描述方法功能>
     * <数据源转换>
     *
     * @param dataSources
     * @return java.util.List<com.it.smn.divide.loadbalance.WeightDataSource>
     * @Author sunmannian
     * @Description
     * @Date 2020/8/31 17:42
     **/
    private List<WeightDataSource> convert(List<DataSource> dataSources) {
        return dataSources.stream().filter(d -> d instanceof AbstractDataSourceWrapper).map(d -> {
            WeightDataSource w = new WeightDataSource();
            w.setCurrentWeight(0);
            w.setOriginalWeight(((AbstractDataSourceWrapper) d).getWeight());
            w.setDataSource(d);
            return w;
        }).collect(Collectors.toList());
    }

    private WeightDataSource smooth(List<WeightDataSource> dataSources, WeightDataSource best) {
        int total = 0;
        for (WeightDataSource dataSource : dataSources) {
            if (dataSource == null) {
                continue;
            }
            total += dataSource.getOriginalWeight();
            dataSource.setCurrentWeight(dataSource.getOriginalWeight() + dataSource.getCurrentWeight());
            if (best == null || dataSource.getCurrentWeight() > best.getCurrentWeight()) {
                best = dataSource;
            }
        }
        if (best == null) {
            return null;
        }
        best.setCurrentWeight(best.getCurrentWeight() - total);
        return best;
    }

    public static void main(String[] args) {
        RoundRobinLoadBalance r = new RoundRobinLoadBalance();
        List<DataSource> dataSources = new ArrayList<>();
        AbstractDataSourceWrapper d1 = new ReadableDataSource(null);
        AbstractDataSourceWrapper d2 = new ReadableDataSource(null);
        AbstractDataSourceWrapper d3 = new ReadableDataSource(null);
        d1.setWeight(1);
        d2.setWeight(5);
        d3.setWeight(1);
        dataSources.add(d3);
        dataSources.add(d1);
        dataSources.add(d2);
        List<WeightDataSource> convert = r.convert(dataSources);
        WeightDataSource weightDataSource = null;
        DataSource select = r.select(dataSources);
        select = r.select(dataSources);
        select = r.select(dataSources);
        select = r.select(dataSources);
        select = r.select(dataSources);
        select = r.select(dataSources);
        select = r.select(dataSources);
        select = r.select(dataSources);
        select = r.select(dataSources);
        select = r.select(dataSources);
    }
}

class WeightDataSource {

    private int originalWeight;

    private int currentWeight;

    private DataSource dataSource;

    public int getOriginalWeight() {
        return originalWeight;
    }

    public void setOriginalWeight(int originalWeight) {
        this.originalWeight = originalWeight;
    }

    public int getCurrentWeight() {
        return currentWeight;
    }

    public void setCurrentWeight(int currentWeight) {
        this.currentWeight = currentWeight;
    }

    public DataSource getDataSource() {
        return dataSource;
    }

    public void setDataSource(DataSource dataSource) {
        this.dataSource = dataSource;
    }
}
